import abc
import sys
import warnings
from typing import Callable, Generic, Optional, TypeVar

from paramiko.client import SSHClient
from zkl_serialization import JsonValue

from jobs_agency.platforms.ssh import SSHAttachedJob, SSHConfig, SSHDetachedJob, SSHPlan, SSHPlatform, ssh_execute
from jobs_agency.skeletons import FakeStore, SkeletonAttachedJob, SkeletonDetachedJob, SkeletonPlatform, StoppedJob

II = TypeVar('II', bound=JsonValue)


class CloudAttachedJob(SkeletonAttachedJob, Generic[II]):
    def __init__(self,
        platform: 'CloudPlatform',
        instance_id: II,
        job_id: str,
        ssh_job: SSHAttachedJob,
        release_script: str | None = None,
    ):
        super().__init__()
        self._platform = platform
        self._instance_id = instance_id
        self._job_id = job_id
        self._ssh_job = ssh_job
        self._release_script = release_script

    @property
    def instance_id(self) -> II:
        return self._instance_id

    @property
    def job_id(self) -> str:
        return self._job_id

    @property
    def ssh_job(self) -> SSHAttachedJob:
        return self._ssh_job

    @property
    def ssh_config(self) -> SSHConfig:
        return self._ssh_job.ssh_config

    @property
    def ssh_client(self) -> SSHClient:
        return self._ssh_job.ssh_client

    @property
    def release_script(self) -> str | None:
        return self._release_script

    def _detach(self):
        if self._ssh_job:
            self._ssh_job.detach()
            self._ssh_job = None
        if self._platform:
            self._platform = None

    def _terminate(self):
        if self._ssh_job:
            if self._release_script:
                print("Release script from config...", file=sys.stderr)
                ssh_execute(self._ssh_job.ssh_client, self._release_script)
            self._ssh_job.terminate()
            self._ssh_job = None
        if self._platform:
            self._platform._release_instance(self._instance_id)
            self._platform._store.remove_job(self._job_id)
            self._platform = None


class CloudDetachedJob(SkeletonDetachedJob, Generic[II]):
    def __init__(self,
        id: str, platform: 'CloudPlatform',
        instance_id: II, ssh_job: Optional[SSHDetachedJob],
    ):
        super().__init__(id, platform)
        self._instance_id = instance_id
        self._ssh_job = ssh_job

    def _check_active(self) -> bool:
        return self._platform._check_instance_running(self._instance_id)

    def _monitor(self):
        if self._ssh_job:
            self._ssh_job.monitor()
        else:
            print("Nothing to monitor: lack of ssh connection.", file=sys.stderr)

    def _terminate(self):
        if self._ssh_job:
            self._ssh_job.terminate()
        self._platform._release_instance(self._instance_id)


class CloudInstanceConfig(abc.ABC):
    @property
    @abc.abstractmethod
    def prepare_func(self) -> Callable[[CloudAttachedJob], None] | None:
        pass

    @property
    @abc.abstractmethod
    def prepare_script(self) -> str | None:
        pass

    @property
    @abc.abstractmethod
    def release_script(self) -> str | None:
        pass


C = TypeVar('C', bound=CloudInstanceConfig)


class CloudPlatform(Generic[II, C], SkeletonPlatform[C, SSHPlan, CloudAttachedJob, CloudDetachedJob], abc.ABC):
    def _new_attached_job(self, config: C) -> CloudAttachedJob:
        try:
            instance_id = self._allocate_instance(config)
            print(f"Created instance: instance_id={instance_id}", file=sys.stderr)

            job_essence = {"instance_id": instance_id}
            job_id = self._store.save_job(None, job_essence)
        except Exception:
            warnings.warn(
                "!!!Attention!!!\n"
                "An error occurred when creating the instance.\n"
                "The instance may have been created and started COSTING MONEY!\n"
                "Due to the error, I have no means to terminate it automatically.\n"
                "Do check it out manually at the cloud service provider side!!")
            raise

        try:
            ssh_config = self._configure_ssh(config, instance_id)
            print(f"Configured SSH: {ssh_config.url}", file=sys.stderr)
            ssh_job = SSHPlatform(FakeStore()).new_attached_job(ssh_config)
            job = CloudAttachedJob(self, instance_id, job_id, ssh_job, config.release_script)

            if config.prepare_func:
                print("Prepare func from config...", file=sys.stderr)
                config.prepare_func(job)
            if config.prepare_script:
                print("Prepare script from config...", file=sys.stderr)
                ssh_execute(ssh_job.ssh_client, config.prepare_script)
            if config.prepare_func or config.prepare_script:
                print("Preparation done!\n", file=sys.stderr)

        except Exception:
            self._release_instance(instance_id)
            raise

        return job

    def _new_detached_job(self, config: C, plan: SSHPlan) -> CloudDetachedJob:
        with self._new_attached_job(config) as attached_job:
            detached_job = self._new_detached_job_from_attached_job(attached_job, plan)
            attached_job.detach()
            return detached_job

    def _new_detached_job_from_attached_job(self, attached_job: CloudAttachedJob, plan: SSHPlan) -> CloudDetachedJob:
        plan = self._configure_plan(attached_job, plan)

        store = FakeStore()
        ssh_job = SSHPlatform(store)._new_detached_job_from_attached_job(attached_job.ssh_job, plan)
        ssh_job_essence = store.job_essence

        job_essence = self._store.load_job(attached_job.job_id)
        job_essence.update({'ssh': ssh_job_essence})
        self._store.save_job(attached_job.job_id, job_essence)
        return CloudDetachedJob(attached_job.job_id, self, attached_job.instance_id, ssh_job)

    def _restore_detached_job(self, job_id: str, job_essence: JsonValue) -> SkeletonDetachedJob | StoppedJob:
        instance_id = job_essence['instance_id']
        ssh_job_essence = job_essence.get('ssh')

        if ssh_job_essence:
            store = FakeStore(ssh_job_essence)
            ssh_job = SSHPlatform(store).get_detached_job("")
        else:
            ssh_job = None
        return CloudDetachedJob(job_id, self, instance_id, ssh_job)

    @abc.abstractmethod
    def _allocate_instance(self, config: C) -> II:
        pass

    @abc.abstractmethod
    def _configure_ssh(self, config: C, instance_id: II) -> SSHConfig:
        pass

    @abc.abstractmethod
    def _configure_plan(self, attached_job: CloudAttachedJob, plan: SSHPlan) -> SSHPlan:
        pass

    @abc.abstractmethod
    def _check_instance_running(self, instance_id: II) -> bool:
        pass

    @abc.abstractmethod
    def _release_instance(self, instance_id: II):
        pass
